﻿#region Copyright information
// 
// Copyright © 2005-2013 Yongkee Cho. All rights reserved.
// 
// This code is a part of the Biological Object Library and governed under the terms of the
// GNU Lesser General  Public License (LGPL) version 2.1 which accompanies this distribution.
// For more information on the LGPL, please visit http://bol.codeplex.com/license.
// 
// - Filename: MultinomialDistribution.cs
// - Author: Yongkee Cho
// - Email: yongkeecho@gmail.com
// - Date Created: 2012-09-06 11:39 AM
// - Last Modified: 2013-01-25 3:59 PM
// 
#endregion
using System;
using System.Collections.Generic;
using System.Linq;
using BOL.Linq;
using BOL.Linq.Probability;
using BOL.Maths.Functions;

namespace BOL.Maths.Distributions
{
    /// <summary>
    /// Represents a multinomial distribution with n trials and probability vector p.
    /// <see>Mathematica: MultivariateStatistics/ref/MultinomialDistribution</see>
    /// <see>http://en.wikipedia.org/wiki/Multinomial_distribution</see>
    /// <typeparam name="TKey"></typeparam>
    /// </summary>
    public sealed class MultinomialDistribution<TKey> : NomialDistributionBase<TKey>, IMultivariateDistribution<int>, IEquatable<MultinomialDistribution<TKey>>
    {
        #region Private variables

        private int _n; // the number of trials

        #endregion

        #region Public properties

        /// <summary>Gets or sets number of trials of a multinomial distribution.</summary>
        public int N { get { return _n; } set { _n = value; } }

        /// <summary>Gets the range of a multinomial distribution.</summary>
        public IRange<int> Domain { get { return new Range<int>(0, _n); } }

        /// <summary>Gets the mean of a multinomial distribution.</summary>
        public double[] Mean { get { return ProbabilityTable.Select(x => x.Value * _n).ToArray(); } }

        /// <summary>Gets the median of a multinomial distribution.</summary>
        public double[] Median { get { throw new Exception("Median is undefined for multinomial distribution."); } }

        /// <summary>Gets the mode of a multinomial distribution.</summary>
        public double[] Mode { get { throw new Exception("Mode is undefined for multinomial distribution."); } }

        /// <summary>Gets the variance of a multinomial distribution.</summary>
        public double[] Variance { get { return ProbabilityTable.Select(x => _n * x.Value * (1 - x.Value)).ToArray(); } }

        /// <summary>Gets the covariance of a multinomial distribution.</summary>
        public double[][] Covariance { get { return ProbabilityTable.Select(x => ProbabilityTable.Select(y => x.Key.Equals(y.Key) ? _n * x.Value * (1 - x.Value) : - _n * x.Value * y.Value).ToArray() ).ToArray(); } }

        /// <summary>Gets the skewness of a multinomial distribution.</summary>
        public double Skewness { get { throw new Exception("Skewness is undefined for multinomial distribution."); } }

        /// <summary>Gets the kurtosis of a multinomial distribution.</summary>
        public double Kurtosis { get { throw new Exception("Kurtosis is undefined for multinomial distribution."); } }

        /// <summary>Gets the entropy of a multinomial distribution.</summary>
        public double Entropy { get { throw new Exception("Entropy is undefined for multinomial distribution."); } }

        #endregion

        #region Constructors

        /// <summary>
        /// Instantiates a multinomial distribution.
        /// </summary>
        /// <param name="probabilityTable">Probability table.</param>
        /// <param name="n">The number of trials.</param>
        public MultinomialDistribution(IDictionary<TKey, double> probabilityTable, int n)
            : base(probabilityTable)
        {
            if (n < 0)
                throw new ArgumentOutOfRangeException("n", "n must be a positive integer.");
            _n = n;
        }

        public MultinomialDistribution(ICollection<TKey> keys, int n)
            : this(keys.Zip(Enumerable.Repeat(1.0, keys.Count()), (k, v) => new { k, v }).ToDictionary(x => x.k, x => x.v), n) { }

        #endregion

        #region ICloneable implementation

        public MultinomialDistribution<TKey> Clone()
        {
            return new MultinomialDistribution<TKey>(ProbabilityTable, _n);
        }

        object ICloneable.Clone()
        {
            return Clone();
        }
        
        #endregion

        #region Public methods

        /// <summary>
        /// Returns the multinomial density given value.
        /// </summary>
        /// <param name="value"></param>
        /// <returns>multinomial density</returns>
        public double Pdf(int[] value)
        {
            if (value.Length != K)
                throw new ArgumentException("Dimension mismatches.", "value");

            if (value.Sum() != _n)
                throw new ArgumentException("Sum of count vector should be n.");

            return ProbabilityTable.Zip(value, (p, v) => Math.Pow(p.Value, v)).
                Aggregate(CountingFunctions.Multinomial(value), (pdf, pow) => pdf * pow);
        }

        /// <summary>
        /// Returns the cummulative multinomial density given value.
        /// </summary>
        /// <see cref="http://controls.engin.umich.edu/wiki/index.php/Multinomial_distributions#Cumulative_Distribution_Function"/>
        /// <param name="value"></param>
        /// <returns>cummulative multinomial density</returns>
        public double Cdf(int[] value)
        {
            if (value.Length != K)
                throw new ArgumentException("Dimension mismatches.", "value");

            if (value.Sum() != _n)
                throw new ArgumentException("Sum of count vector should be n.");

            var count = value.Select(x => x).ToArray(); // copy count vector

            double cdf = 0;
            while (count[K - 1] < _n) // sum pdf until n1 = 0, n2 = 0, ..., nk-1 = 0, nk = N
            {
                cdf += Pdf(count); // sum pdf
                for (var i = 0; i < K - 1; i++) // for n1, ... nk-1
                {
                    if (count[i] > 0) // to find ni of ci > 0
                    {
                        count[i]--; // ci = ci - 1;
                        count[K - 1] = _n; // calculate ck = N - c1 - c2 - ... - ck-1;
                        for (var j = 0; j < K - 1; j++)
                            count[K - 1] -= count[j];
                        break; // once a ci value changed, exit loop and update cdf
                    }
                }
            }

            return cdf;
        }

        /// <summary>
        /// Returns the value of the multinomial distribution for given probability.
        /// </summary>
        /// <param name="p">probability</param>
        /// <returns></returns>
        public int[] Quantile(double p)
        {
            throw new NotImplementedException();
        }

        /// <summary>
        /// Returns a random value from the distribution.
        /// <see>http://en.wikipedia.org/wiki/Multinomial_distribution</see>
        /// </summary>
        /// <returns>random value</returns>
        public int[] Sample(Random r)
        {
            var sample = new TKey[_n].Select(x => ProbabilityTable.Sample(r));
            return ProbabilityTable.Select(x => sample.Count(y => y.Equals(x.Key))).ToArray();
        }

        /// <summary>
        /// Estimates parameters of the Bernoulli distribution using maximum likelihood.
        /// </summary>
        public new void MaximumLikelihoodEstimate(IEnumerable<TKey> source)
        {
            if (source == null)
                throw new ArgumentNullException("source");
            var list = source.ToList();
            base.MaximumLikelihoodEstimate(list);
            N = list.Count();
        }

        public void MaximumLikelihoodEstimate(IEnumerable<int[]> source)
        {
            throw new NotSupportedException();
        }

        /// <summary>
        /// Estimates parameters of the Bernoulli distribution using Bayesian methods.
        /// <param name="priorParameters">Dirichlet priors</param>
        /// </summary>
        public new void BayesianEstimate(IEnumerable<TKey> source, params dynamic[] priorParameters)
        {
            if (source == null)
                throw new ArgumentNullException("source");
            var list = source.ToList();
            base.BayesianEstimate(list);
            N = list.Count();
        }

        public void BayesianEstimate(IEnumerable<int[]> source, params dynamic[] priorParameters)
        {
            throw new NotSupportedException();
        }

        #endregion

        #region IEquatable<MultinomialDistribution<TKey>> implementation

        public bool Equals(MultinomialDistribution<TKey> other)
        {
            return _n.Equals(other._n) & ProbabilityTable.DictionaryEqual(other.ProbabilityTable);
        }

        #endregion

        #region Object overriden

        public override int GetHashCode()
        {
            return ProbabilityTable.Aggregate(1, (hashCode, p) => hashCode ^ p.Value.GetHashCode()) ^ N.GetHashCode();
        }

        public override bool Equals(object other)
        {
            if (other == null)
                throw new ArgumentNullException("other");

            if (!(other is MultinomialDistribution<TKey>))
                throw new InvalidCastException("The 'other' argument is not a MultinomialDistribution object.");
            
            return Equals(other as MultinomialDistribution<TKey>);
        }

        public override string ToString()
        {
            return String.Format("Multinomial(n = {0}, {1})", _n, String.Join(", ", ProbabilityTable.Select(p => String.Format("p_{0} = {1:0.0000}", p.Key, p.Value))));
        }

        #endregion
    }
}
